import os
import json
import shutil
import logging
from pathlib import Path
from typing import List, Dict, Any

from source.RFT.utils import logger, ensure_dir, write_json_file
from source.RFT.models import ScenarioMetadata, TraceSample, DatasetMetadata

def create_output_structure(
    output_dir: str,
    input_dir: str,
    scenarios: List[ScenarioMetadata],
    samples: List[TraceSample],
    config: Dict[str, Any]
) -> bool:
    """
    创建输出目录结构，复制原始数据文件，生成样本记录JSON文件
    
    参数:
        output_dir: 输出目录路径
        input_dir: 输入目录路径
        scenarios: 场景元数据列表
        samples: 样本记录列表
        config: 配置信息
        
    返回:
        操作是否成功
    """
    output_path = ensure_dir(output_dir)
    logger.info(f"开始生成输出目录结构: {output_path}")
    
    # 1. 创建关卡文件夹
    for scenario in scenarios:
        level_path = output_path / scenario.level_folder
        ensure_dir(level_path)
        
        # 源文件路径
        scenario_folder = Path(input_dir) / scenario.folder_name
        source_cache_file = scenario_folder / "cache" / "telemetry_loader.pkl"
        source_similarity_file = scenario_folder / "art_features" / "features" / "similarity_results.csv"
        
        # 目标文件路径
        target_cache_file = level_path / "telemetry_loader.pkl"
        target_similarity_file = level_path / "similarity_results.csv"
        
        # 复制telemetry_loader.pkl
        try:
            shutil.copy2(source_cache_file, target_cache_file)
            logger.debug(f"已复制缓存文件: {source_cache_file} -> {target_cache_file}")
        except Exception as e:
            logger.error(f"复制缓存文件失败: {e}")
            return False
            
        # 复制similarity_results.csv（如果存在）
        if source_similarity_file.exists():
            try:
                shutil.copy2(source_similarity_file, target_similarity_file)
                logger.debug(f"已复制相似度文件: {source_similarity_file} -> {target_similarity_file}")
            except Exception as e:
                logger.warning(f"复制相似度文件失败: {e}")
    
    # 2. 生成样本记录JSON文件
    samples_json_path = output_path / "samples.json"
    
    # 创建元数据对象
    metadata = DatasetMetadata(config)
    metadata.update_from_scenarios(scenarios)
    metadata.update_from_samples(samples)
    
    # 构建JSON数据
    json_data = {
        "metadata": metadata.to_dict(),
        "scenarios": [vars(s) for s in scenarios],
        "samples": [vars(s) for s in samples]
    }
    
    # 写入JSON文件
    if not write_json_file(samples_json_path, json_data):
        logger.error(f"写入样本记录JSON文件失败: {samples_json_path}")
        return False
        
    logger.info(f"输出生成完成: {len(scenarios)}个关卡, {len(samples)}个样本")
    
    return True